import torch
import time
import numpy as np
from importlib import import_module
from utils import build_dataset, build_iterator, get_time_dif
import argparse
from train import train, init_network

parser = argparse.ArgumentParser(description="中文文本分类！")
parser.add_argument('--model', type=str, default='TextCNN')
parser.add_argument('--embedding', type=str, default='pre_trained', help='choose random or pre_trained')
parser.add_argument('--word', type=bool, default=False, help='True for word split, or False for char split')
args = parser.parse_args()


if __name__ == '__main__':
    dataset_path = 'dataset'

    model_name = args.model
    if args.embedding == 'random':
        embedding = ""
    else:
        embedding = "embedding_SougouNews.npz"

    x = import_module('models.' + model_name)
    config = x.Config(dataset_path, embedding)

    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed(1)
    torch.backends.cudnn.deterministic = True

    start_time = time.time()
    print("**********开始加载数据!************")
    vocab, train_data, test_data, dev_data = build_dataset(args.word, config)
    train_iter = build_iterator(train_data, config)
    test_iter = build_iterator(test_data, config)
    dev_iter = build_iterator(dev_data, config)
    time_dif = get_time_dif(start_time)

    print("数据集加载用时：", time_dif)

    config.n_vocab = len(vocab)
    model = x.Model(config).to(config.device)

    def init_weights(model):
        for name, w in model.named_parameters():
            if 'embedding' in name:
                pass
            else:
                if 'weight' in name:
                    torch.nn.init.xavier_normal_(w)
                if 'bias' in name:
                    torch.nn.init.constant_(w, 0)
    init_weights(model)

    # init_network(model)

    train(config, model, train_iter, test_iter, dev_iter)

    print("asd")




